import math
from typing import List, Tuple, Union, Literal
import pygame
import pygame.gfxdraw
import numpy as np

Number = Union[int, float]
Directions = Literal["UP", "DOWN"]
SemaphoreStates = Literal["STOP", "PASS", "WARNING", "DISABLED"]
# Direction
_id = 0


def new_id():
    global _id
    _id += 1
    return _id


class BaseTrackShape:
    def to_dict(self):
        raise NotImplementedError()


class StraightTrack(BaseTrackShape):
    def __init__(
        self, start_pos: Tuple[Number, Number], end_pos: Tuple[Number, Number]
    ) -> None:
        self.start_pos = start_pos
        self.end_pos = end_pos

    def to_dict(self):
        return {
            "type": "StraightTrack",
            "start_pos": self.start_pos,
            "end_pos": self.end_pos,
        }


class TrackData:
    """
    轨道模型。表示一条轨道，包括起点和终点，以及轨道的状态。
    """

    def __init__(
        self,
        id: str,
        shapes: List[StraightTrack],
        semaphore_up_position: Tuple[Number, Number] = None,
        semaphore_down_position: Tuple[Number, Number] = None,
    ) -> None:
        # self.start_pos = start_pos
        self.id = id
        self.shapes = shapes
        self.semaphore_up_state: SemaphoreStates = "PASS"
        self.semaphore_down_state: SemaphoreStates = "PASS"
        start_pos = self.shapes[0].start_pos
        end_pos = self.shapes[-1].end_pos
        if end_pos[0] > start_pos[0]:  # 区间下行方向是向右的
            coe = 1
        else:
            coe = -1
        ORIGINAL_SEMAPHORE_UP_POSITION = (
            start_pos[0] + 10 * coe,
            start_pos[1] - 10,
        )
        ORIGINAL_SEMAPHORE_DOWN_POSITION = (
            end_pos[0] - 10 * coe,
            end_pos[1] - 10,
        )
        self.semaphore_up_position: Tuple[Number, Number] = (
            semaphore_up_position or ORIGINAL_SEMAPHORE_UP_POSITION
        )
        self.semaphore_down_position: Tuple[Number, Number] = (
            semaphore_down_position or ORIGINAL_SEMAPHORE_DOWN_POSITION
        )

    def get_shapes(self):
        return {
            "semaphore_up_position": self.semaphore_up_position,
            "semaphore_down_position": self.semaphore_down_position,
            "shapes": [shape.to_dict() for shape in self.shapes],
        }

    def get_semaphore_color(
        self, semaphore: SemaphoreStates
    ) -> Tuple[Number, Number, Number]:
        if semaphore == "PASS":
            return (0, 255, 0)
        elif semaphore == "STOP":
            return (255, 0, 0)
        else:
            return (255, 255, 0)


class Route:
    """
    进路
    """

    def __init__(
        self,
        name: str,
        route_items: List[Tuple[str, Directions]],
        protected_semaphores: List[
            Tuple[str, Directions]
        ],  # 受保护的信号机，必须置为STOP
    ) -> None:
        self.name = name
        self.route_items = route_items
        self.protected_semaphores = protected_semaphores


class Section:
    """
    轨道区段。包括若干条轨道，里面可以规划若干条进路
    """

    def __init__(self, name: str, tracks: List[TrackData], routes: List[Route]) -> None:
        self.name = name
        self.tracks = tracks
        self.routes = routes

    def get_track_shapes(self):
        return [
            {"name": self.name, "shapes": track.get_shapes()} for track in self.tracks
        ]


class RailwayNetwork:
    """
    铁路网
    """

    def __init__(
        self,
        sections: List[Section],
        # junctions反映的是轨道之间的连接关系，六元组分为前后两部分各3个，
        # 3个数据分别表示轨道的区段名、轨道名和方向
        junctions: List[Tuple[str, str, Directions, str, str, Directions]],
    ) -> None:
        self.sections = sections
        self.junctions = junctions

    def get_shapes(self):
        """
        获取轨道网的形状
        """
        track_shapes = []
        for section in self.sections:
            track_shapes.extend(section.get_track_shapes())
        return {"track_shapes": track_shapes}
